import cv2
import numpy

class PatternRecognizer:
    def __init__(self):
        self.patterns = None
    
    def set_patterns(self, pattern_dict):
        self.patterns = pattern_dict

    def match_img(self, img, template, threshold, nms_threshold=0.5):
        if img is None or template is None:
            raise ValueError("图像或模板路径错误，请检查路径")

        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        template = cv2.cvtColor(template, cv2.COLOR_BGR2GRAY)

        h, w = template.shape[:2]
        halfh = h // 2
        halfw = w // 2
        
        # 使用归一化相关系数匹配方法
        result = cv2.matchTemplate(img, template, cv2.TM_CCOEFF_NORMED)
        
        # 获取所有匹配位置（高于阈值的位置）
        locations = numpy.where(result >= threshold)
        
        rects = []  # 位置矩形坐标
        points = []  # 位置中点
        for pt in zip(*locations[::-1]):
            rects.append([pt[0], pt[1], w, h])
            points.append([pt[0] + halfw, pt[1] + halfh])

        # 应用非极大值抑制 (NMS)
        rects = numpy.array(rects)
        scores = result[locations]
        indices = cv2.dnn.NMSBoxes(
            rects.tolist(), scores.tolist(), threshold, nms_threshold
        )

        # # 绘制结果:
        # if len(indices) > 0:
        #     img_display = img.copy()
        #     for r in rects[indices]:
        #         cv2.rectangle(
        #             img_display,
        #             (r[0], r[1]),  # Left Up
        #             (r[0] + w, r[1] + h),  # Right Down
        #             (0, 0, 255),  # Color
        #             2  # Line width
        #         )
        #     cv2.imshow('Matches', img_display)
        #     cv2.waitKey(0)
        #     cv2.destroyAllWindows()

        points = numpy.array(points)
        if len(indices) > 0:
            points = points[indices]
        return points

    def run(self, img, threshold=0.9):
        res = {}
        if self.patterns is None:
            return res
        for k, v in self.patterns.items():
            template = cv2.imread(v)
            res[k] = self.match_img(img, template, threshold)
        return res
